package cn.itcast.manus.agent.browser;

import cn.hutool.core.convert.Convert;
import cn.hutool.core.io.resource.ResourceUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import cn.itcast.manus.agent.prompt.PromptManagement;
import cn.itcast.manus.constants.Constant;
import cn.itcast.manus.service.PageContentExtractService;
import com.microsoft.playwright.BrowserContext;
import com.microsoft.playwright.BrowserContext.WaitForPageOptions;
import com.microsoft.playwright.Page;
import com.microsoft.playwright.Page.ScreenshotOptions;
import com.microsoft.playwright.Page.WaitForLoadStateOptions;
import com.microsoft.playwright.Page.WaitForSelectorOptions;
import com.microsoft.playwright.Page.WaitForURLOptions;
import com.microsoft.playwright.TimeoutError;
import com.microsoft.playwright.options.LoadState;
import jakarta.annotation.Resource;
import lombok.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;

import java.util.*;
import java.util.Map.Entry;
import java.util.function.Predicate;
import java.util.stream.Stream;

/**
 * 浏览器页面会话管理类
 * 封装了Playwright的页面操作功能，用于管理页面的导航、元素交互、内容提取和DOM分析等操作
 * 作为manus调用浏览器操作的核心组件，提供了丰富的浏览器自动化方法
 */
@Slf4j
public class PageSession {

    // JavaScript脚本：移除页面中的高亮容器元素
    private static final String PLAYWRIGHT_CONTAINER_REMOVE = "document.getElementById(\"playwright-highlight-container\")?.remove()";

    // JavaScript脚本：移除页面中所有元素的data-testid属性
    private static final String PLAYWRIGHT_DATATESTID_REMOVE = """
            document.querySelectorAll('*').forEach(element => {
                    if (element.hasAttribute('data-testid')) {
                      element.removeAttribute('data-testid');
                    }
                  });
            """;

    // 常量：表示操作数量的参数名
    public static final String AMOUNT = "amount";

    // 常量：表示高亮索引的键名
    private static final String HIGHLIGHT_INDEX = "highlightIndex";

    // 等待超时时间，单位为毫秒
    private static final long WAIT_TIMEOUT = 5000L;

    // 提示管理服务，用于获取系统提示模板
    @Resource
    private PromptManagement promptManagement;

    // 页面内容提取服务，用于从页面中提取指定内容
    @Resource
    private PageContentExtractService pageContentExtractService;

    // 存储页面所有元素的映射，键为元素ID，值为元素的JSON对象表示
    private final Map<Integer, JSONObject> allElement = new TreeMap<>();

    // DOM树的根节点ID
    private Integer rootId;

    // 构建DOM树的JavaScript代码
    private final String buildDomJS;

    // 要收集的HTML元素属性集合
    private final Set<String> attributeKey;

    // Playwright浏览器上下文对象
    private final BrowserContext browserContext;

    // 存储操作过程中打开的页面列表
    private final List<Page> openInOperation = new LinkedList<>();

    // 当前页面索引
    @Setter
    @Getter
    private int pageIndex;

    // 页面列表信息 当前浏览器全部tab信息
    @Getter
    private String pageListInfo;

    // 导航标志：1表示发生了导航，0表示未发生导航
    private int flagNavi = 0;

    // 新页面标志：1表示打开了新页面，0表示未打开新页面
    private int flagNewPage = 0;

    // 标签页变更标志：1表示标签页已变更，0表示标签页未变更
    private int flagTabChanged = 0;

    // 原始页面内容
    private String originalContent;

    /**
     * 构造函数，初始化PageSession实例
     * @param browserContext Playwright浏览器上下文对象
     */
    public PageSession(BrowserContext browserContext) {
        this.browserContext = browserContext;
        this.pageIndex = 0;
        // 初始化需要收集的HTML元素属性集合
        this.attributeKey = Set.of("title", "type", "name", "role", "tabindex", "aria-label", "placeholder", "value", "alt", "aria-expanded", "class");
        // 加载构建DOM树的JavaScript代码
        this.buildDomJS = ResourceUtil.readUtf8Str("js/buildDomTree.js");
    }

    /**
     * 清除页面状态标志
     * 重置导航、新页面和标签页变更的状态标志
     */
    public void clearFlag() {
        this.flagNavi = 0;
        this.flagNewPage = 0;
        this.flagTabChanged = 0;
    }

    /**
     * 刷新页面信息，包括当前页面的标题、URL等
     * 更新pageListInfo和originalContent字段
     */
    private void refreshPageInfo() {
        var pageListInfo = new StringBuilder();
        var page = listPageWithoutExtension();
        var pageArr = page.toArray(Page[]::new);
        // 创建浏览器标签页信息列表 将浏览器页面数组转换成包含页面信息的对象列表
        var lst = Stream.iterate(0, i -> i + 1)//创建一个从0开始，每次递增1的无限流
                .limit(pageArr.length) //限制流的大小为页面数组的长度
                .map(i -> new BrowserTabInfo(i, pageArr[i].title(), pageArr[i].url()))//: Lambda表达式，接收索引i，创建一个新的BrowserTabInfo对象
                .toList();
        pageListInfo.append(JSONUtil.toJsonStr(lst)).append("\n");
        pageListInfo.append("current at index:").append(this.pageIndex).append("\n");
        pageListInfo.append("using 'switchTab' action to change browser tab\n");
        this.pageListInfo = pageListInfo.toString();
        this.originalContent = currentPage().content();
    }

    /**
     * 检查页面状态是否发生显著变化
     * 在浏览器自动化测试和操作中，页面状态的变化检测非常重要：
     *
     * 页面变化感知：通过检查特定元素是否存在，可以判断页面是否发生了显著变化
     * 稳定性保障：设置超时时间可以避免程序因等待页面变化而无限挂起
     * 自动化流程控制：根据页面状态变化决定后续操作，实现更智能的自动化流程
     * @return 如果页面状态发生显著变化则返回true，否则返回false
     */
    public boolean isStatusSignificantChanged() {
        // 检查是否发生了导航、新页面打开或标签页变更
        if (flagNavi != 0 || flagNewPage != 0 || flagTabChanged != 0) {
            return true;
        }
        try {
            var page = currentPage();
            //创建并配置等待页面选择器的选项
            var opt = new WaitForSelectorOptions();
            opt.setTimeout(WAIT_TIMEOUT);

            var handle = page.waitForSelector("#playwright-highlight-container", opt);
            return handle == null;
        } catch (Exception e) {
            log.trace("significantChanged by exception", e);
            return true;
        }
    }

    /**
     * 重新分析页面，刷新页面信息并构建DOM树
     * 清除旧的分析容器，执行DOM树构建脚本，并解析结果
     */
    private void reAnalyze() {
        this.refreshPageInfo();
        this.removeAnalyzeContainer();
        var page = currentPage();

        // 执行构建DOM树的JavaScript脚本
        var result = page.evaluate(buildDomJS);
        var jo = JSONUtil.parseObj(result);
        var map = jo.getJSONObject("map");
        // 获取根节点ID
        this.rootId = jo.getInt("rootId");
        // 将所有的节点信息写入到allElement中
        map.forEach((k, v) -> allElement.put(Convert.toInt(k), (JSONObject) v));
    }

    /**
     * 等待新页面的出现，超时时间为5秒
     * @return 如果成功等待到新页面，返回包含新页面的Optional对象；否则返回空Optional
     */
    private Optional<Page> waitForNewPage() {
        try {
            var opt = new WaitForPageOptions();
            opt.setTimeout(WAIT_TIMEOUT);
            var newPage = this.browserContext.waitForPage(opt, () -> {});
            // 等待新页面加载完成
            this.waitPageIdle(newPage);
            this.flagNewPage = 1;
            return Optional.of(newPage);
        } catch (TimeoutError e) {
            log.trace("no page found after waiting 5 s");
            return Optional.empty();
        }
    }

    /**
     * 等待页面加载完成，超时时间为5秒
     * @param p 需要等待加载的页面对象
     */
    private void waitPageIdle(Page p) {
        try {
            // 等待页面加载状态完成
            p.waitForLoadState();
            var opt = new WaitForLoadStateOptions();
            opt.setTimeout(WAIT_TIMEOUT);
            // 等待网络空闲状态
            p.waitForLoadState(LoadState.NETWORKIDLE, opt);
        } catch (TimeoutError e) {
            log.trace("ignore waiting loadstatus timeout.", e);
        }
    }

    /**
     * 获取当前页面对象
     * @return 当前页面对象
     */
    private Page currentPage() {
        return this.listPageWithoutExtension().get(pageIndex);
    }

    /**
     * 列出所有非扩展页面
     * @return 非扩展页面的列表
     */
    private List<Page> listPageWithoutExtension() {
        return browserContext.pages().stream()
                .filter(p -> !p.url().toLowerCase().startsWith("chrome-extension"))
                .toList();
    }

    /**
     * 移除页面中的分析容器和测试ID属性
     * 清理页面中可能影响后续操作的元素和属性
     */
    public void removeAnalyzeContainer() {
        var page = this.listPageWithoutExtension();
        page.forEach(p -> {
            p.evaluate(PLAYWRIGHT_CONTAINER_REMOVE);
            p.evaluate(PLAYWRIGHT_DATATESTID_REMOVE);
        });
    }

    /**
     * 清理页面资源，关闭打开的页面
     * 在会话结束时调用，释放浏览器资源
     */
    public void cleanUp() {
        this.removeAnalyzeContainer();
        this.openInOperation.forEach(Page::close);
    }

    /**
     * 等待页面导航完成
     * @param page 要等待的页面对象
     * @param oldUrl 导航前的URL
     * @return 如果导航成功完成返回true，否则返回false
     */
    private boolean waitNavi(Page page, String oldUrl) {
        try {
            var waitForURLOptions = new WaitForURLOptions();
            waitForURLOptions.setTimeout(WAIT_TIMEOUT);
            Predicate<String> urlChanged = s -> !s.equalsIgnoreCase(oldUrl);
            page.waitForURL(urlChanged, waitForURLOptions);
            this.flagNavi = 1;
            return true;
        } catch (TimeoutError e) {
            log.trace("ignore waiting loadstatus timeout.", e);
            return false;
        }
    }

    /**
     * 定位到指定页面，并更新页面索引
     * @param page 要定位的页面对象
     */
    private void locateOn(Page page) {
        this.pageIndex = listPageWithoutExtension().indexOf(page);
        this.openInOperation.add(page);
    }

    /**
     * 向上滚动页面
     * @param arg 包含滚动参数的映射，其中AMOUNT键对应的值表示滚动的像素数
     * @return 操作结果消息
     */
    public String scrollUp(Map<String, Object> arg) {
        var page = currentPage();
        int amount = Convert.toInt(arg.get(AMOUNT));
        page.evaluate(String.format("window.scrollBy(0, -%s);", amount));
        return "SUCCESS";
    }

    /**
     * 向下滚动页面
     * @param arg 包含滚动参数的映射，其中AMOUNT键对应的值表示滚动的像素数
     * @return 操作结果消息
     */
    public String scrollDown(Map<String, Object> arg) {
        var page = currentPage();
        int amount = Convert.toInt(arg.get(AMOUNT));
        page.evaluate(String.format("window.scrollBy(0, %s);", amount));
        return "SUCCESS";
    }

    /**
     * 获取当前页面的状态信息，包括URL、页面数据、滚动信息等
     * @param currentStep 当前步骤
     * @param maxStep 最大步骤
     * @return 格式化的页面状态信息
     */
    public String getCurrentPageStatus(int currentStep, int maxStep) {
        this.reAnalyze();
        var page = this.currentPage();
        var params = new HashMap<>();
        params.put("url", page.url());
        params.put("pageListInfo", getPageListInfo());
        // 生成页面元素树数据
        var pageData = fromAllElementByTree(rootId);
        params.put("pageData", pageData);
        params.put("stepData", String.format("%s/%s", currentStep, maxStep));
        params.put("dateTime", String.valueOf(new Date()));

        // 获取滚动信息
        var sc = this.getScrollInfo();
        params.put("scrollData",
                String.format(
                        "...%s pixels above,%s pixels below,using extract_content action to see more information...",
                        sc.getPixelsAbove(), sc.getPixelsBelow()));

        // 使用页面状态提示模板格式化结果
        return StrUtil.format(promptManagement.getPrompt(Constant.Prompts.PAGE_STATUS), params);
    }

    /**
     * 检查指定索引的节点是否包含高亮标记
     * @param indexId 节点索引ID
     * @return 如果节点或其子节点包含高亮标记则返回true，否则返回false
     */
    private boolean containsHighlighted(Integer indexId) {
        var node = this.allElement.get(indexId);
        if (node.containsKey(HIGHLIGHT_INDEX)) {
            return true;
        }
        return childIdxList(node).stream().anyMatch(this::containsHighlighted);
    }

    /**
     * 获取节点的子节点索引列表
     * @param node 节点对象
     * @return 子节点索引列表
     */
    private List<Integer> childIdxList(JSONObject node) {
        return Optional.of(node)
                .filter(j -> j.containsKey("children"))
                .map(j -> j.getJSONArray("children")
                        .stream()
                        .map(Convert::toInt)
                        .toList())
                .orElse(Collections.emptyList());
    }

    /**
     * 递归生成页面元素树的字符串表示
     * @param indexId 起始节点索引ID
     * @return 格式化的页面元素树字符串
     */
    private String fromAllElementByTree(Integer indexId) {
        var node = this.allElement.get(indexId);
        var childList = this.childIdxList(node);
        // 递归生成子节点内容
        var childContent = childList.stream()
                .map(this::fromAllElementByTree)
                .filter(StrUtil::isNotBlank)
                .reduce((x, y) -> x + "\n" + y)
                .orElse(StrUtil.EMPTY);

        // 处理不可见节点
        if (!node.getBool("isVisible", false)) {
            return childContent;
        }

        // 处理文本节点
        if ("TEXT_NODE".equals(node.getStr("type"))) {
            return node.getStr("text") + "\n" + childContent;
        }

        // 处理带有高亮标记的节点
        if (node.containsKey(HIGHLIGHT_INDEX)) {
            var highlightIndex = node.getStr(HIGHLIGHT_INDEX);
            var tag = Optional.ofNullable(node.getStr("tagName")).orElse("?");
            var attr = Optional.ofNullable(node.getJSONObject("attributes")).map(this::collectAttribute)
                    .filter(s -> !s.equals(childContent)).orElse(StrUtil.EMPTY);
            var childHighLighted = childList.stream().anyMatch(this::containsHighlighted);
            return childHighLighted ? childContent : String.format("[%s]<%s %s;%s/>", highlightIndex, tag, attr, StrUtil.removeAllSuffix(childContent, "\n"));
        }
        return childContent;
    }

    /**
     * 收集节点的属性信息，只收集预定义的属性集合中的属性
     * @param node 包含属性的节点对象
     * @return 格式化的属性字符串
     */
    private String collectAttribute(JSONObject node) {
        return node.entrySet().stream()
                .filter(e -> attributeKey.contains(e.getKey()))
                .map(Entry::getValue)
                .map(String::valueOf)
                .reduce((x, y) -> x + ";" + y)
                .orElse(StrUtil.EMPTY);
    }

    /**
     * 获取页面滚动信息
     * @return 包含滚动信息的ScrollInfo对象
     */
    public ScrollInfo getScrollInfo() {
        var page = currentPage();
        var scrollY = Convert.toInt(page.evaluate("window.scrollY"));
        var viewportHeight = Convert.toInt(page.evaluate("window.innerHeight"));
        var totalHeight = Convert.toInt(page.evaluate("document.documentElement.scrollHeight"));
        return ScrollInfo.builder()
                .pixelsAbove(scrollY)
                .pixelsBelow(totalHeight - (scrollY + viewportHeight))
                .build();
    }

    /**
     * 根据目标描述提取相关的内容数据
     * @param goal 要提取的内容描述
     * @return 提取的内容
     */
    @Tool(description = "根据目标，提取相关的内容数据")
    public String extractContent(@ToolParam(description = "要提取的内容描述") String goal) {
        return this.pageContentExtractService.extractContent(originalContent, this.fromAllElementByTree(rootId), goal);
    }

    /**
     * 切换到指定的页面标签页
     * @param pageId 页面ID
     * @return 操作结果消息
     */
    @Tool(description = "切换到指定的页面标签页")
    public String switchTab(@ToolParam(description = "页面id") int pageId) {
        this.pageIndex = pageId;
        this.reAnalyze();
        this.flagTabChanged = 1;
        return "Action Result:page index change to " + pageId;
    }

    /**
     * 点击指定索引的元素，并处理可能的新页面打开或导航
     * @param index 元素下标
     * @return 操作结果消息
     */
    @Tool(description = "点击指定索引的元素")
    public String clickElement(@ToolParam(description = "下标") int index) {
        Page page = currentPage();
        String oldUrl = page.url();
        String testId = String.format("zp-%s", index);
        // 通过测试ID定位并点击元素
        page.getByTestId(testId).click();

        // 处理点击后可能出现的新页面或导航
        String postfix = waitForNewPage().map(p -> {
            locateOn(p);
            return "new page found after click action:" + p.title();
        }).orElseGet(() -> {
            if (waitNavi(page, oldUrl)) {
                return "navigated to:" + page.title();
            }
            return StrUtil.EMPTY;
        });
        return "Action result: Element clicked!" + postfix;
    }

    /**
     * 导航到指定URL
     * @param url URL地址
     * @return 操作结果消息
     */
    @Tool(description = "导航到指定的url")
    public String goToUrl(@ToolParam(description = "url地址") String url) {
        Page page = browserContext.newPage();
        page.navigate(url);
        waitNavi(page, StrUtil.EMPTY);
        locateOn(page);
        return "Action result: Navigated to " + url;
    }

    /**
     * 在指定索引的输入框中输入文本
     * @param index 输入框下标
     * @param text 要输入的文本内容
     * @return 操作结果消息
     */
    @Tool(description = "输入框中输入文本")
    public String inputText(@ToolParam(description = "下标") int index, @ToolParam(description = "文本内容") String text) {
        Page page = currentPage();
        String testId = String.format("zp-%s", index);
        var locator = page.getByTestId(testId);
        locator.clear();
        locator.fill(text);
        return "Action result: input text success";
    }

    /**
     * 等待指定时间
     * @param seconds 等待时间（秒）
     * @return 操作结果消息
     */
    @Tool(description = "等待指定时间")
    public String wait(@ToolParam(description = "等待时间(秒)") int seconds) {
        try {
            Thread.sleep(seconds * 1_000L);
            return "Action result: you have wait " + seconds + " second(s)";
        } catch (InterruptedException e) {
            return "Action result: InterruptedException(" + e.getMessage() + ")";
        }
    }

    /**
     * 用于存储页面滚动信息的内部类
     */
    @Data
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class ScrollInfo {
        // 视口上方的像素数
        private int pixelsAbove;
        // 视口下方的像素数
        private int pixelsBelow;
    }

    /**
     * 用于存储浏览器标签页信息的内部类
     */
    @Data
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class BrowserTabInfo {
        // 标签页索引
        private int index;
        // 标签页标题
        private String title;
        // 标签页URL
        private String url;
    }
}
